#!/bin/bash
###############################################################################
#
#  EGSnrc script to run individual pbsdsh tasks
#  Copyright (C) 2020 National Research Council Canada
#
#  This file is part of EGSnrc.
#
#  EGSnrc is free software: you can redistribute it and/or modify it under
#  the terms of the GNU Affero General Public License as published by the
#  Free Software Foundation, either version 3 of the License, or (at your
#  option) any later version.
#
#  EGSnrc is distributed in the hope that it will be useful, but WITHOUT ANY
#  WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
#  FOR A PARTICULAR PURPOSE.  See the GNU Affero General Public License for
#  more details.
#
#  You should have received a copy of the GNU Affero General Public License
#  along with EGSnrc. If not, see <http://www.gnu.org/licenses/>.
#
###############################################################################
#
#  Author:          Frederic Tessier, 2020
#
#  Contributors:
#
###############################################################################
#
#  This script is not meant to be called directly, but rather from the script
#  egs-parallel-pbsdsh (PBS distributed shell)
#
###############################################################################


### help function
function help {
    log "HELP"
    cat <<EOF

    usage:

        $(basename $0) pbsdsh_dir basename nthread first delay 'command'

    arguments:

        pbsdsh_dir  existing directory to save task files for job numbers
        basename    simulation input file name, without ".egsinp" extension
        nthread     number of task to run under dsh
        first       first job index
        delay       user-specified delay in seconds between individual jobs
        command     command to run, in quotes

    note:

        This script is not meant to be called directly, but rather from the
        egs-parallel-pbsdsh script

EOF
}

### timestamp function
function timestamp {
    printf "EGSnrc egs-parallel $(date -u "+%Y-%m-%d (UTC) %H:%M:%S.%N")"
}

### log function to write messages to log file and standard output
function log {
    msg="$(timestamp): $1\n"
    printf "$msg"
}

### quit function for errors, with source, line, message and command
function quit {
    lineno=$1
    msg=$2
    case $3 in
        help)  cmd="help";;
        *)     cmd="";;
    esac
    log "$(basename $0): line $lineno: $msg"; $cmd; log "QUIT."; exit 1
}

### quit function if simulation is done
function quit_if_done {
    if [ -r $basename.egsjob ]; then
        done=$(grep -o END $basename.egsjob)
        if [ "$done" = "END" ]; then
            log "$jobstr: QUIT (simulation already finished)"
            exit
        fi
    fi
}

### go to pbs working directory
cd $PBS_O_WORKDIR

### parse command-line arguments (simplistic)
args_min=6
if [ "$#" -lt $args_min ]; then
    quit $LINENO "only $# arguments provided; at least $args_min required" help
fi
pbsdsh_dir=$1
basename=$2
nthread=$3
first=$4
delay=$5
command=$6

### task label and task file
task=${PBS_TASKNUM}
taskstr="task $task"
prefix="$pbsdsh_dir/${basename}_"
taskfile=${prefix}${task}.task
log "$taskstr: host $HOSTNAME: $taskfile"
touch $taskfile

### wait until all tasks have launched
delta=2
filecount=$(ls -Ub1 -- $prefix*.task | wc -l)
while [ $filecount -lt $nthread ]; do
    log "$taskstr: wait $delta seconds for all tasks to start ($filecount/$nthread)"
    sleep $delta
    filecount=$(ls -Ub1 -- $prefix*.task | wc -l)
done

### select manager job (first one in file list)
manager=$(ls $prefix*.task | sed -n 1p)           # using sed instead of head to curb broken pipe error
manager=$(basename ${manager#${prefix}} .task)

### manager assigns task indices
if [ $manager == $task ]; then
    log "$taskstr: manager: assigning $nthread tasks"
    job=$first
    for f in $(ls ${prefix}*.task); do
        printf "$job\n" > $f
        job=$((job+1))
    done
fi

### wait for my job index (which is printed in task file by manager job)
delta=2
while [ ! -s $taskfile ]; do
    log "$taskstr: wait $delta seconds for job index"
    sleep $delta
done

### job index and label
job=$(cat $taskfile)
jobstr=$(printf "job %04d" $job)
log "$jobstr <- $taskstr"

# log the host and pid of this job
log "$jobstr: host=$(hostname)"
log "$jobstr: BEGIN pid=$$"

### manage jobs to avoid bottleneck and race conditions
if [ $job -eq 1 ]; then

    # log host and pid of job 1 in .egsjob file
    log "$jobstr: BEGIN host=$(hostname) pid=$$" > $basename.egsjob

else

    # all jobs wait a fixed delay (relative to first job)
    delta=1
    log "$jobstr: wait $delta seconds (initial delay)"
    sleep $delta

    # wait until there is an .egsjob file (maximum 120 seconds)
    total=0
    delta=10
    limit=120
    while [ ! -e $basename.egsjob ]; do

        # quit if simulation is already done
        quit_if_done

        # otherwise wait for egsjob file
        log "$jobstr: wait $delta seconds (no $basename.egsjob file after $total seconds)"
        sleep $delta
        total=$((total+$delta))
        if [ $total -gt $limit ]; then
            log "$jobstr: QUIT (no $basename.egsjob file after $limit seconds)"
            exit
        fi
    done

    # quit if simulation is already done
    quit_if_done

    # offset all jobs by a fixed delay (relative to previous job)
    delta=100000
    log "$jobstr: wait $((job*$delta)) microseconds (default job offset delay)"
    for j in $(seq 1 $job); do
        usleep $delta
        quit_if_done
    done

    # extra user-specified delay between each job
    delta=$delay
    if [ $delta -gt 0 ]; then
        log "$jobstr: wait $((job*$delta)) seconds (user job offset delay)"
        for j in $(seq 1 $job); do
            sleep $delta
            quit_if_done
        done
    fi

    # report on lock file content
    if [ -r $basename.lock ]; then
        content=$(cat $basename.lock)
        log "$jobstr: found $basename.lock: $content"
    fi

fi

### launch the job if simulation is not done already
quit_if_done
log "$jobstr: RUN $command -b -P $nthread -j $job -f $first"
$command -b -P $nthread -j $job -f $first

### report that this job is done
log "$jobstr: DONE."

### report that the simulation (job 1) is done
if [ $job -eq 1 ]; then
    log "$jobstr: END host=$(hostname) pid=$$" >> $basename.egsjob
fi